from cv2 import cv2
import os

class VidSrc(object):
    def __init__(self, root_path, cps=1, clip_suffix='.avi', nsec_per_clip=3600, fps=1):
        self.version = '20210215'
        print('[VidSrc.__init__] version = %s' % self.version)
        self.root_path = root_path
        self.cps = cps
        self.fps = fps
        self.step = None
        self.play_spf = float(1. / cps)
        self.clip_flist = []
        self.nfr_list = []
        self.__list_clips(clip_suffix)
        self.curr_clip_idx = 0
        self.curr_fr_idx = 0
        self.nsec_per_clip = nsec_per_clip
        self.nfr_per_clip = nsec_per_clip * fps
        self.nfr = 0
        self.img_buffer = []
        self.h_vid = None
        self.__init_nfr_list()
        self.__connect_clip(0)

    def __list_clips(self, suffix='.avi'):    
        """
            列出所有avi文件
        """
        for f in os.listdir(self.root_path):
            if f.endswith(suffix):
                self.clip_flist.append(f)
        self.clip_flist.sort()

    def __init_nfr_list(self):
        for f_clip in self.clip_flist:
            fp_clip = os.path.join(self.root_path, f_clip)
            self.h_vid = cv2.VideoCapture(fp_clip)
            curr_nfr = self.h_vid.get(cv2.CAP_PROP_FRAME_COUNT)
            self.nfr_list.append(curr_nfr)
            self.nfr += curr_nfr

    def __connect_clip(self, clip_idx):
        """
            连接到指定视频文件
            clip_idx: 视频序号
            分类处理以下情况：
                return True 成功连接指定视频
                return False 超出视频列表范围 or 到达所有视频末尾
        """
        if clip_idx < len(self.clip_flist): # 在视频列表范围内
            fp_clip = os.path.join(self.root_path, self.clip_flist[clip_idx]) # 获取视频绝对路径
            #if isinstance(self.h_vid, cv2.VideoCapture) and self.h_vid.isOpened():
            #     self.h_vid.release()
            self.curr_clip_idx = clip_idx # 保存当前视频 index
            self.h_vid = cv2.VideoCapture(fp_clip) # 获取当前视频句柄
            curr_nfr = self.h_vid.get(cv2.CAP_PROP_FRAME_COUNT) # 当前视频总帧数
            if curr_nfr < self.nfr_per_clip: # 视频帧数不足
                print('[VidSrc.warning] insufficient nfr %d < standard %d' % (curr_nfr, self.nfr_per_clip))
            self.fps = self.h_vid.get(cv2.CAP_PROP_FPS) # 获取视频实际fps
            self.step = int(self.fps / self.cps) # 计算帧间间隔
            if self.fps % self.cps != 0 or self.cps > self.fps: # cps不能整除fps
                print('[VidSrc.warning] invalid cps = %d setting for video fps = %d, fps : cps should be integar.' % (self.cps, self.fps))
            else: # fps是cps的整数倍
                print('[VidSrc.info] frame step = %d' % (self.step))
            self.curr_fr_idx = 0
            return True
        elif clip_idx == len(self.clip_flist): # 到达视频末尾
            print('[VidSrc.info] reach the end of all clips.')
            return False
        else: # 超出视频列表范围
            print('[VidSrc.warning] clip index is out of range.')
            return False

    def jump_to_fr_inclip(self, fr_idx):    
        """
            当前视频片段内，跳转到指定帧
            fr_idx: 指定帧序号
            分类处理以下情况：
                return True 成功跳转到指定帧
                return True 跳转到下一个视频
                return False 到达所有视频末尾 or 远超出当前视频片段的帧范围
        """
        if fr_idx <= self.nfr_per_clip:
            diff = fr_idx - self.curr_fr_idx
            if diff < 0 or diff > 1:
                self.h_vid.set(cv2.CAP_PROP_POS_FRAMES, fr_idx)
            self.curr_fr_idx = fr_idx
            return True
        else:
            if fr_idx > self.nfr_per_clip + self.fps:
                print('[VidSrc.warning] frame index is far beyond current clip.')
                return False
            print('[VidSrc.info] jump to next clip.')
            return self.jump_to_next_clip()

    def jump_to_next_clip(self):
        self.curr_clip_idx += 1
        return self.__connect_clip(self.curr_clip_idx)

    def jump_to_fr(self, clip_idx, fr_idx):
        """
            跳转到指定视频片段的指定帧
            clip_idx: 视频序号
            fr_idx: 指定帧序号
            分类处理以下情况：
                return True 成功跳转
                return True 跳转到下一个视频
                return False 到达所有视频末尾 or 超出视频列表范围
                return False fr_idx远超出当前视频片段的帧范围
        """
        if clip_idx != self.curr_clip_idx:
            is_valid_clip = self.__connect_clip(clip_idx)
            if not is_valid_clip: # 跳转视频失败
                return False
        if fr_idx < 0:
            n_fr_inclip = int(self.h_vid.get(cv2.CAP_PROP_FRAME_COUNT))
            fr_idx += n_fr_inclip
        is_valid_fr = self.jump_to_fr_inclip(fr_idx)
        if not is_valid_fr: # 跳转帧失败
            return False
        return True
    
    def jump_to_sec(self, clip_idx, sec_idx):
        self.jump_to_fr(clip_idx, sec_idx * self.step)

    def imread(self, mode='bgr'):
        if isinstance(self.h_vid, cv2.VideoCapture) and self.h_vid.isOpened():
            _, img = self.h_vid.read()
            if img is None:
                return None
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) if mode == 'rgb' else img
            fr_idx = self.curr_fr_idx + self.step
            is_finish = self.jump_to_fr_inclip(fr_idx)
            return img if is_finish else None
        else:
            print('[VidSrc.warning] invalid video handler.')
            return None
    
    def get_clip_idx(self):
        return self.curr_clip_idx

    def get_inclip_fr_idx(self):
        return int(self.h_vid.get(cv2.CAP_PROP_POS_FRAMES))

    def get_abs_fr_idx(self):
        relative_fr_idx = self.get_inclip_fr_idx()
        prev_clip_nfr = 0
        for i in range(self.curr_clip_idx):
            prev_clip_nfr += self.nfr_list[i]
        return prev_clip_nfr + relative_fr_idx

    def get_curr_nfr(self):
        return int(self.h_vid.get(cv2.CAP_PROP_FRAME_COUNT))

    def get_nfr(self):
        return self.nfr

    def get_nsec(self):
        return int(self.nfr / self.fps)

    def is_end(self):
        # 检查是否为最后一个视频片段
        return self.curr_clip_idx > len(self.clip_flist) - 1 and self.curr_fr_idx > self.get_curr_nfr() - self.fps

def create_flysleep_video(root_path, fps=1, nsec_per_clip=3600): # 果蝇睡眠追踪
    return VidSrc(root_path, cps=1, clip_suffix='.avi', nsec_per_clip=nsec_per_clip, fps=fps)

def create_mousesleep_video(root_path, fps=1, nsec_per_clip=3600): # 小鼠睡眠追踪
    return VidSrc(root_path, cps=1, clip_suffix='.avi', nsec_per_clip=nsec_per_clip, fps=fps)

def create_rat_pair_video(root_path, fps=25, nsec_per_clip=600): # 大鼠配对行为分析
    return VidSrc(root_path, cps=fps, clip_suffix='.avi', nsec_per_clip=nsec_per_clip, fps=fps)

def test_vid_src():
    root_path = '/home/youngway/data/mtrace/20201202_c57-males-dvr1/1'
    # h_vid = create_rat_pair_video(root_path, fps=25, nsec_per_clip=600) # 创建视频句柄
    h_vid = create_mousesleep_video(root_path) # 创建视频句柄
    cv2.namedWindow("trace") # 打开新窗口
    i_sec = 0.
    while(1):
        img = h_vid.imread()
        i_sec += h_vid.play_spf
        cv2.putText(img, '#%d fr, #%.2f sec' % (h_vid.curr_fr_idx, i_sec), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, color=[0, 255, 0])
        cmd = cv2.waitKey(100)
        if img is None or cmd == 27:
            break
        elif cmd == ord('b'):
            h_vid.jump_to_fr(h_vid.curr_clip_idx, h_vid.curr_fr_idx - 100)
            print(h_vid.get_inclip_fr_idx())
        else:
            cv2.imshow("test", img)

if __name__ == '__main__':
    test_vid_src()